
#include "egb_segment.h"

bool operator<(const edge &a, const edge &b)
{
  return a.w < b.w;
}

// random color
rgb random_rgb(){ 
  rgb c;

  c.r = (uchar) rand()%255;
  c.g = (uchar) rand()%255;
  c.b = (uchar) rand()%255;

  return c;
}

//===============================================================

universe::universe(int elements)
{
	elts = new uni_elt[elements];
	
	num = elements;

	for (int i = 0; i < elements; i++)
	{
		elts[i].rank = 0;
		elts[i].size = 1;
		elts[i].p = i;
	}
}
  
universe::~universe()
{
  delete [] elts;
}

int universe::find(int x)
{
	int y = x;

	while (y != elts[y].p)
		y = elts[y].p;

	elts[x].p = y;

	return y;
}

void universe::join(int x, int y)
{
	if (elts[x].rank > elts[y].rank)
	{
		elts[y].p = x;
		elts[x].size += elts[y].size;
	}
	else
	{
		elts[x].p = y;
		elts[y].size += elts[x].size;
		if (elts[x].rank == elts[y].rank)
		elts[y].rank++;
	}

	num--;
}

//===============================================================

egl_segment::egl_segment()
{
	sigma = 0.8f;
	k = 300.0f;
	min_size = 512;

	srand((unsigned int)time(NULL));
}

float egl_segment::diff(Mat &img, int x1, int y1, int x2, int y2)
{
	float result;

	if(img.channels()==1)
	{
		result = sqrt( pow( (float)img.at<unsigned char>(y1,x1)-img.at<unsigned char>(y2,x2), 2) );
	}
	else if(img.channels()==3)
	{
		result = sqrt( pow( (float)img.at<unsigned char>(y1,x1*3  )-img.at<unsigned char>(y2,x2*3  ), 2) 
					 + pow( (float)img.at<unsigned char>(y1,x1*3+1)-img.at<unsigned char>(y2,x2*3+1), 2)
					 + pow( (float)img.at<unsigned char>(y1,x1*3+2)-img.at<unsigned char>(y2,x2*3+2), 2));
	}

	return result;
}

/*
 * Segment a graph
 *
 * Returns a disjoint-set forest representing the segmentation.
 *
 * num_vertices: number of vertices in graph.
 * num_edges: number of edges in graph
 * edges: array of edges.
 * c: constant for treshold function.
 */
universe* egl_segment::segment_graph(int num_vertices, int num_edges, edge *edges, float c)
{
	clock_t t0=clock();

	// sort edges by weight
	std::sort(edges, edges + num_edges);

	clock_t t1=clock();
	printf ( "s1 %.2f\n", (float)t1-t0 );

	// make a disjoint-set forest
	universe *u = new universe(num_vertices);

	// init thresholds
	float *threshold = new float[num_vertices];
	for (int i = 0; i < num_vertices; i++)
		threshold[i] = c;

	clock_t t2=clock();
	printf ( "s2 %.2f\n", (float)t2-t1 );

	// for each edge, in non-decreasing weight order...
	for (int i = 0; i < num_edges; i++)
	{
		edge *pedge = &edges[i];
    
		// components conected by this edge
		int a = u->find(pedge->a);
		int b = u->find(pedge->b);
		if (a != b)
		{
			if ((pedge->w <= threshold[a]) && (pedge->w <= threshold[b]))
			{
				u->join(a, b);
				a = u->find(a);
				threshold[a] = pedge->w + (c / u->size(a));
			}
		}
	}

	clock_t t3=clock();
	printf ( "s3 %.2f\n", (float)t3-t2 );

	// free up
	delete threshold;
	return u;
}

int egl_segment::process(Mat &src, Mat &dst)
{
	clock_t t0=clock();

	GaussianBlur(src, dst, Size(3,3), sigma);
	src.release();

	clock_t t1=clock();
	printf ( "t1 %.2f\n", (float)t1-t0 );

	int height = dst.rows;
	int width = dst.cols;

	edge *edges = new edge[dst.rows * dst.cols*4];	//edge_types = 4

	int num = 0;
	for (int y = 0; y < dst.rows; y++)
	{
		for (int x = 0; x < dst.cols; x++)
		{
			if (x < width-1)
			{
				edges[num].a = y * width + x;
				edges[num].b = y * width + (x+1);
				edges[num].w = diff(dst, x, y, x+1, y);
				num++;
			}

			if (y < height-1)
			{
				edges[num].a = y * width + x;
				edges[num].b = (y+1) * width + x;
				edges[num].w = diff(dst, x, y, x, y+1);
				num++;
			}

			if ((x < width-1) && (y < height-1))
			{
				edges[num].a = y * width + x;
				edges[num].b = (y+1) * width + (x+1);
				edges[num].w = diff(dst, x, y, x+1, y+1);
				num++;
			}

			if ((x < width-1) && (y > 0))
			{
				edges[num].a = y * width + x;
				edges[num].b = (y-1) * width + (x+1);
				edges[num].w = diff(dst, x, y, x+1, y-1);
				num++;
			}
		}
	}

	clock_t t2=clock();
	printf ( "t2 %.2f\n", (float)t2-t1 );

	universe *u = segment_graph(width*height, num, edges, k);

	clock_t t3=clock();
	printf ( "t3 %.2f\n", (float)t3-t2 );

	// post process small components
	for (int i = 0; i < num; i++)
	{
		int a = u->find(edges[i].a);
		int b = u->find(edges[i].b);
		if ((a != b) && ((u->size(a) < min_size) || (u->size(b) < min_size)))
		u->join(a, b);
	}

	delete [] edges;

	clock_t t4=clock();
	printf ( "t4 %.2f\n", (float)t4-t3 );

	num_ccs = u->num_sets();

	// pick random colors for each component
	if(dst.channels()==1)
	{
		unsigned char* colors = new unsigned char[width*height];
		for (int i = 0; i < width*height; i++)
		{
			colors[i] = rand()%255;
		}
  
		for (int y = 0; y < height; y++)
		{
			for (int x = 0; x < width; x++)
			{
				int comp = u->find(y * width + x);
				
				dst.at<unsigned char>(y, x) = colors[comp];
			}
		}

		delete [] colors;  
	}
	else if(dst.channels()==3)
	{
		rgb *colors = new rgb[width*height];
		for (int i = 0; i < width*height; i++)
		{
			colors[i] = random_rgb();
		}
  
		for (int y = 0; y < height; y++)
		{
			for (int x = 0; x < width; x++)
			{
				int comp = u->find(y * width + x);

				dst.at<unsigned char>(y, x*3)   = colors[comp].r;
				dst.at<unsigned char>(y, x*3+1) = colors[comp].g;
				dst.at<unsigned char>(y, x*3+2) = colors[comp].b;
			}
		}

		delete [] colors;  
	}

	delete u;

	clock_t t5=clock();
	printf ( "t5 %.2f\n", (float)t5-t4 );

	return 0;
}